import os
import re
from hex_string_utils import split_hex_file  # 修改为新文件名的导入
from CLittleEdianToBigEdian import convert_little_endian_to_big_endian
from convert_b_matrix import convert_b_matrix_row_to_col_major

# 解析文件路径中的mnk信息，返回(m, n, k)
def parse_mnk_from_path(path):
    match = re.search(r'm(\d+)n(\d+)k(\d+)', path)
    if match:
        m, n, k = map(int, match.groups())
        return m, n, k
    else:
        return None, None, None

# 生成大小端转换后的文件名
def get_big_endian_filename(file_path):
    dirname, basename = os.path.split(file_path)
    # 针对c_int32_m..._hex.txt格式
    match = re.match(r'(c_int32_)(m\d+n\d+k\d+)(_hex)?\.txt', basename)
    if match:
        new_basename = f"{match.group(1)}big_endian_{match.group(2)}_hex.txt"
        return os.path.join(dirname, new_basename)
    # 其他类型直接在.txt前加_big_endian
    return file_path.replace('.txt', '_big_endian.txt')

# 生成B矩阵列优先文件名
def get_col_major_filename(file_path):
    # 在.txt前加_col_major_hex
    return file_path.replace('.txt', '_col_major_hex.txt')

# 主批量处理函数
def batch_process_fp_dirs(fp_dirs):
    for fp_dir in fp_dirs:
        # 判断是fp16还是fp32
        if 'fp16' in fp_dir:
            split_len = 4
            byte_width = 2
        elif 'fp32' in fp_dir:
            split_len = 8
            byte_width = 4
        else:
            print(f"未知类型目录: {fp_dir}")
            continue
        # 遍历所有txt文件
        for subdir, _, files in os.walk(fp_dir):
            for file in files:
                if not file.endswith('.txt'):
                    continue
                file_path = os.path.join(subdir, file)
                print(f"\n处理文件: {file_path}")
                # 解析mnk
                m, n, k = parse_mnk_from_path(file_path)
                # 1. 分割，直接覆盖原文件
                split_hex_file(file_path, file_path, split_len)
                print(f"分割完成: {file_path}")
                # 2. 大小端转换，生成新文件名
                be_file = get_big_endian_filename(file_path)
                convert_little_endian_to_big_endian(file_path, be_file, byte_width)
                print(f"大小端转换完成: {be_file}")
                # 3. B矩阵行优先转列优先，生成新文件名
                if os.path.basename(file).startswith('b_') and n and k:
                    col_major_file = get_col_major_filename(file_path)
                    convert_b_matrix_row_to_col_major(be_file, col_major_file, rows=k, cols=n)
                    print(f"B矩阵列优先转换完成: {col_major_file}")

if __name__ == "__main__":
    # 需要处理的目录
    fp_dirs = [
        r"D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/fp32",
        r"D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/fp16"
    ]
    batch_process_fp_dirs(fp_dirs)
    print("\n全部处理完成！") 